from fastapi import FastAPI
from pydantic import BaseModel
from notionai import NotionAI
import json
from newbing.bing_request import init_conversation_with_magic, getStruct, get_anser
app = FastAPI()

dict_mode = '_new'

notion_pre_list = ["Continue write",
"Summarize",
"Improve writing",
"Fix spelling grammar",
"Explain this",
"Make longer",
"Make shorter",
"Find action items",
"Simplify language",
"Brainstorm ideas",
"Blog post",
"Outline",
"Social media post",
"Press release",
"Creative story",
"Poem",
"Meeting agenda",
"Pros and cons list",
"Job description",
"Sales email",
"Recruiting email"]
trans_pre_list = [
"english",
"korean",
"chinese",
"japanese",
"spanish",
"russiab",
"french",
"german",
"italian",
"protuguese",
"dutch",
"indonesia",
"tagalog",
"vietnamese"
]
def ai_go(token,space_id,prompt,context):
    ai = NotionAI(token,space_id)
    if prompt not in notion_pre_list and prompt not in trans_pre_list:
        res = ai.help_me_write(prompt,context)
    else:
        if prompt in trans_pre_list:
            res = ai.translate(prompt,context)
        else:
            if prompt == "Continue write":
                res = ai.continue_write(context)
            if prompt == "Summarize":
                res = ai.summarize(context)
            if prompt == "Improve writing":
                res = ai.improve_writing(context)
            if prompt == "Fix spelling grammar":
                res = ai.fix_spelling_grammar(context)
            if prompt == "Explain this":
                res = ai.explain_this(context)
            if prompt == "Make longer":
                res = ai.make_longer(context)
            if prompt == "Make shorter":
                res = ai.make_shorter(context)
            if prompt == "Find action items":
                res = ai.find_action_items(context)
            if prompt == "Simplify language":
                res = ai.simplify_language(context)
            if prompt == "Brainstorm ideas":
                res = ai.brainstorm_ideas(context)
            if prompt == "Outline":
                res = ai.outline(context)
            if prompt == "Social media post":
                res = ai.social_media_post(context)
            if prompt == "Press release":
                res = ai.press_release(context)
            if prompt == "Creative story":
                res = ai.creative_story(context)
            if prompt == "Poem":
                res = ai.poem(context)
            if prompt == "Meeting agenda":
                res = ai.meeting_agenda(context)
            if prompt == "Pros and cons list":
                res = ai.pros_cons_list(context)
            if prompt == "Job description":
                res = ai.job_description(context)
            if prompt == "Sales email":
                res = ai.sales_email(context)
            if prompt == "Recruiting email":
                res = ai.recruiting_email(context)
    return res

with open('./plus_key.json','r') as f:
    plus_dict = json.load(f)

class Item(BaseModel):
    TOKEN: str
    SPACE_ID: str
    prompt: str
    context: str
    human_history: list = []
    assistant_history: list = []
        
class TOKENItem(BaseModel):
    TOKEN: str

class NewBingCreationConversationItem(BaseModel):
    MagicURL: str
    cookies: str

class NewBingChatItem(BaseModel):
    conversation_signature: str
    client_id: str
    conversation_id: str
    invocation_id: int
    prompt: str
    conversation_style:  str
    wss_link: str
        
def chat(token,space_id,prompt,context,hh,ah):
    ai = NotionAI(token,space_id)
    return ai.chat(prompt,hh,ah,context)
           
@app.post("/write/")
async def write(item: Item):
    if len(item.human_history) == 1:
        res = ai_go(item.TOKEN, item.SPACE_ID,item.prompt,context=item.context)
    else:
        res = chat(item.TOKEN, item.SPACE_ID,item.prompt,context=item.context,hh=item.human_history,ah=item.assistant_history)
    return res

@app.post("/spaceid/")
async def write(item: TOKENItem):
    ai = NotionAI(item.TOKEN,space_id='none')
    res = ai.get_spaces(item.TOKEN)
    return res

@app.post("/NewBingCreationConversation/")
async def create_conversation(item: NewBingCreationConversationItem):
    conversation_signature, client_id, conversation_id = init_conversation_with_magic(cookies=item.cookies,magic=item.MagicURL)
    print("(conversation_signature, client_id, conversation_id)")
    print((conversation_signature, client_id, conversation_id))
    res_struct = {
            "conversation_signature": conversation_signature,
            "client_id": client_id,
            "conversation_id": conversation_id,
            "debugInf": 'success!',
        }
    return res_struct

@app.post("/NewBingChat/")
async def run_bing_chat(item: NewBingChatItem):
    request = getStruct(item.conversation_signature, 
                        item.client_id, 
                        item.conversation_id, 
                        item.invocation_id,          # NOTE: invocation_id is should be updated each time
                        item.prompt,
                        item.conversation_style)
    res = (
            await get_anser(request,item.wss_link)
    )["item"]["messages"][1]["adaptiveCards"][0]["body"][0]["text"]
    return res

